import { error } from 'cc';
import { Input, Session } from '../../ONNX/WONNX/WONNXCore';
import { DataType, Model, ModelMetadataExtensions } from '../Model';
import { Tensor, TensorShape } from '../Tensor';
import { ModelAnalyzer } from './ModelAnalyzer';

export abstract class IWorker {
	abstract prepareForInput(inputShapes: { [name: string]: TensorShape }, dataType?: DataType): void;
	abstract setInput(x: Tensor): void;
	abstract setInput(name: string, x: Tensor): void;
	abstract execute(): Promise<void>;
	abstract execute(input: Tensor): Promise<void>;
	abstract execute(inputs: { [name: string]: Tensor }): Promise<void>;
	abstract startManualSchedule(): void;
	abstract startManualSchedule(input: Tensor): void;
	abstract startManualSchedule(inputs: { [name: string]: Tensor }): void;
	abstract flushSchedule(blocking?: boolean): void;
	readonly scheduleProgress: number = 0;
	abstract peekOutput(): Tensor;
	abstract peekOutput(name: string): Tensor;
	abstract peekConstants(layerName: string): Tensor[];
	abstract summary(): string;

	dispose(): void {
		// do nothing
	}
}

export enum Device {
	GPU = 1 << 8,
	CPU = 1 << 9,
	Auto = 1 << 15,
}

export enum WorkerType {
	Compute = 1 | Device.GPU,
	TS = 2 | Device.CPU,
	Auto = 0 | Device.Auto,
}

export interface WorkerConfiguration {
	verbose: boolean;
	compareAgainstType: WorkerType;
	compareLogLevel?: number;
	compareEpsilon?: number;
	takeoverWeights?: boolean;
}

export class WorkerFactory {
	static createWorker(type: WorkerType, model: Model, modelBytes: ArrayBuffer, verbose: boolean): IWorker {
		const workerConfiguration = {
			verbose: verbose,
			compareAgainstType: type,
			compareLogLevel: 0,
			compareEpsilon: 0,
			takeoverWeights: false,
		};

		return new GenericWorker(model, modelBytes, workerConfiguration);
	}
}

export class GenericWorker implements IWorker {
	private _inputShapes: { [name: string]: TensorShape } = {};
	private _defaultInputName: string = '';
	private _session: Session | null = null;
	private _outputData: { [name: string]: number[] } = {};
	private _input: Input = null!;
	private _model: Model = null!;
	private _defaultOutputName: string = '';
	private _verbose: boolean = false;
	private _workerType: WorkerType = WorkerType.Auto;
	private _activationDataType: DataType = DataType.Float32;

	constructor(model: Model, modelBytes: ArrayBuffer, workerConfiguration: WorkerConfiguration) {
		this._model = model;
		Session.fromBytes(new Uint8Array(modelBytes)).then((session) => {
			this._session = session;
		}).catch(err => {
			error(err);
		});

		this._defaultInputName = ModelAnalyzer.getDefaultInputName(model);
		this._defaultOutputName = ModelAnalyzer.getDefaultOutputName(model);

		this._verbose = workerConfiguration.verbose;
		this._workerType = workerConfiguration.compareAgainstType;
	}

	prepareForInput(inputShapes: { [name: string]: TensorShape }, dataType?: DataType | undefined): void {
		this._inputShapes = inputShapes;
		this._activationDataType = dataType || DataType.Float32;
	}
	setInput(x: Tensor): void;
	setInput(name: string, x: Tensor): void;
	setInput(name: Tensor | string, x?: Tensor): void {
		if (this._model.inputs.length <= 1 && name == this._defaultInputName) {
			this.prepareForInput({ name: x!.shape }, this._activationDataType);
		}

		if (this._input == null) {
			this._input = new Input();
		}

		if (name instanceof Tensor && x === undefined) {
			this._inputShapes[this._defaultInputName] = name.shape;
			this._input.insert(name.name, new Float32Array(name.cache));
		} else if (typeof name === 'string') {
			this._inputShapes[name] = x!.shape;
			this._input.insert(name, new Float32Array(x!.cache));
		}
	}
	async execute(): Promise<void>;
	async execute(input: Tensor): Promise<void>;
	async execute(inputs: { [name: string]: Tensor }): Promise<void>;
	async execute(inputs?: Tensor | { [name: string]: Tensor }): Promise<void> {
		if (inputs instanceof Tensor) {
			this.setInput(inputs);
		} else {
			for (const name in inputs) {
				this.setInput(name, inputs[name]);
			}
		}

		if (this._session) {
			this._outputData = await this._session.run(this._input);
		}

		this._input.free();
		this._input = null!;
	}
	startManualSchedule(): void;
	startManualSchedule(input: Tensor): void;
	startManualSchedule(inputs: { [name: string]: Tensor }): void;
	startManualSchedule(inputs?: unknown): void {
		throw new Error('Method not implemented.');
	}
	flushSchedule(blocking?: boolean | undefined): void {
		throw new Error('Method not implemented.');
	}
	scheduleProgress: number = 0;
	peekOutput(): Tensor;
	peekOutput(name: string): Tensor;
	peekOutput(name?: string): Tensor {
		if (name) {
			const tensor = ModelMetadataExtensions.getTensorByName(this._model, name);
			tensor.cache = Array.from(this._outputData[name]);

			return tensor;
		} else {
			const tensor = ModelMetadataExtensions.getTensorByName(this._model, this._defaultOutputName);
			tensor.name = this._defaultOutputName;
			tensor.cache = Array.from(this._outputData[this._defaultOutputName]);
			return tensor;
		}
	}
	peekConstants(layerName: string): Tensor[] {
		throw new Error('Method not implemented.');
	}
	summary(): string {
		throw new Error('Method not implemented.');
	}

	dispose(): void {
		this._session?.free();
	}
}
